open Core_kernel.Std
open Printf

type port_direction = Input | Output
type port = int * port_direction

type binary_op =
  | And
  | Or
  | Xor
  | Add

let string_of_binary_op = function
  | And -> "&"
  | Or -> "|"
  | Xor -> "^"
  | Add -> "+"

type unary_op =
  | Not

let string_of_unary_op = function
  | Not -> "~"

type expr = 
  | Sym of string
  | Sym_part of string * (int * int)
  | Literal of int * int (* width, value *)
  | Cond of expr * expr * expr
  | Binary of binary_op * expr * expr
  | Unary of unary_op * expr
  | Concat of expr list

type vlog_instance = {
  class_name : string;
  parameters : (string * int) list;
  ports : (string * port * expr) list;
}

let rec string_of_expr = function
  | Sym s -> s
  | Sym_part (s, (lo, hi)) ->
      if hi-lo > 1
      then sprintf "%s[%d:%d]" s (hi-1) lo
      else sprintf "%s[%d]" s lo
  | Literal (width, value) ->
      sprintf "%d'd%d" width value
  | Cond (cond, t, f) ->
      sprintf "(%s?%s:%s)"
        (string_of_expr cond)
        (string_of_expr t)
        (string_of_expr f)
  | Binary (op, a, b) ->
      sprintf "(%s%s%s)"
        (string_of_expr a)
        (string_of_binary_op op)
        (string_of_expr b)
  | Unary (op, e) ->
      sprintf "(%s%s)" (string_of_unary_op op) (string_of_expr e)
  | Concat exprs ->
      sprintf "{%s}"
        (String.concat ~sep:"," (List.map exprs ~f:string_of_expr))

let rec make_binary_expr (op : binary_op) (exprs : expr list) : expr =
  match exprs with
  | [] -> assert false
  | [expr] -> expr
  | expr :: exprs' -> Binary (op, expr, make_binary_expr op exprs')

type assign = (string * (int * int)) * expr
type reg = int * expr * expr option (* width, next_value, write_enable_opt *)

type vlog_module = {
  port_alist : (string * port) list;
  wire_alist : (string * int) list;
  assigns : assign list; 
  instance_alist : (string * vlog_instance) list;
  reg_alist : (string * reg) list;
  has_clock : bool;
}

let vlog_of_named_port (name, (width, direction)) : string =
  let dir_s =
    match direction with
    | Input -> "input"
    | Output -> "output"
  in
  if width > 1
  then sprintf "%s [%d:0] %s" dir_s (width-1) name
  else sprintf "%s %s" dir_s name

type sym_kind =
  | Sym_input
  | Sym_output
  | Sym_internal
  | Sym_reg

type sym_info = {
  width : int;
  kind : sym_kind;
  drivers : string list array;
}

let find_net (net_map : sym_info String.Map.t) (net_name : string) : sym_info =
  try String.Map.find_exn net_map net_name
  with Not_found ->
    failwithf "no such net: %s" net_name ()

let check_width
  ((a : expr), (a_width : int))
  ((b : expr), (b_width : int))
  : unit =
  if a_width <> b_width
  then begin
    failwithf "expression width mismatch: %s (%d) and %s (%d)"
      (string_of_expr a) a_width (string_of_expr b) b_width ()
  end

let rec check_expr (net_map : sym_info String.Map.t) (expr : expr) : expr * int =
  let check = check_expr net_map in
  match expr with
  | Sym net_name ->
      let sym_info = find_net net_map net_name in
      expr, sym_info.width
  | Sym_part (net_name, (lo, hi)) ->
      assert (lo < hi);
      let sym_info = find_net net_map net_name in
      let expr' =
        if lo = 0 && hi = sym_info.width
        then Sym net_name
        else expr
      in
      if hi > sym_info.width
      then begin
        failwithf "part select out of bounds: %s (%d)"
          (string_of_expr expr) sym_info.width ()
      end;
      expr', hi-lo
  | Literal (width, _) ->
      expr, width
  | Cond (c, t, f) ->
      let c', _ = check c in
      let t', t_width = check t in
      let f', f_width = check f in
      check_width (t, t_width) (f, f_width);
      Cond (c', t', f'), t_width
  | Binary (op, a, b) ->
      let a', a_width = check a in
      let b', b_width = check b in
      check_width (a, a_width) (b, b_width);
      Binary (op, a', b'), a_width
  | Unary (op, e) ->
      let e', e_width = check e in
      Unary (op, e'), e_width
  | Concat exprs ->
      let exprs', widths = List.unzip (List.map exprs ~f:check) in
      let width = List.fold widths ~init:0 ~f:(+) in
      Concat exprs', width

let write_named_vlog_module ((name : string), (m : vlog_module)) (filepath : string) : unit =
  let net_info_alist : (string * sym_info) list =
    let port_net_alist =
      m.port_alist
      |> List.map ~f:begin fun (name, (width, direction)) ->
        let kind =
          match direction with
          | Input -> Sym_input
          | Output -> Sym_output
        in
        let drivers = Array.create [] ~len:width in
        name, { width; kind; drivers }
      end
    in
    let wire_net_alist =
      m.wire_alist
      |> List.map ~f:begin fun (name, width) ->
        let drivers = Array.create [] ~len:width in
        name, { width; kind = Sym_internal; drivers }
      end
    in
    let reg_net_alist =
      List.map m.reg_alist ~f:begin fun (name, (width, _, _)) ->
        name, { width; kind = Sym_reg; drivers = [||] }
      end
    in
    List.concat [ port_net_alist; wire_net_alist; reg_net_alist ]
  in
  let net_map : sym_info String.Map.t =
    String.Map.of_alist_exn net_info_alist
  in
  fprintf stderr "writing file %s\n" filepath;
  let out_chan = open_out filepath in
  fprintf out_chan "module %s(\n  %s\n);\n" name
    begin
      let ports =
        if m.has_clock
        then ("clk", (1, Input)) :: m.port_alist
        else m.port_alist
      in
      List.map ports ~f:vlog_of_named_port
      |> String.concat ~sep:",\n  ";
    end;
  let print_object object_type name width =
    if width > 1
    then fprintf out_chan "%s [%d:0] %s;\n" object_type (width-1) name
    else fprintf out_chan "%s %s;\n" object_type name
  in
  List.iter m.wire_alist ~f:begin fun (name, width) ->
    print_object "wire" name width
  end;
  if m.has_clock
  then begin
    m.reg_alist
    |> List.iter ~f:(fun (name, (width, _, _)) -> print_object "reg" name width)
  end;
  let add_driver sym_info i driver =
    sym_info.drivers.(i) <- driver :: sym_info.drivers.(i)
  in
  List.iter m.assigns ~f: begin fun ((lhs_net_name, (lo, hi)), expr) ->
    let expr, rhs_width = check_expr net_map expr in
    let expr_s = string_of_expr expr in
    let sym_info = find_net net_map lhs_net_name in
    begin match sym_info.kind with
    | Sym_input -> failwithf "assign to input port %s" lhs_net_name ()
    | Sym_reg -> failwithf "assign to reg %s" lhs_net_name ()
    | _ -> ()
    end;
    let lhs_width = hi-lo in
    check_width (Sym lhs_net_name, lhs_width) (expr, rhs_width);
    if lo = 0 && hi = sym_info.width
    then fprintf out_chan "assign %s = %s;\n" lhs_net_name expr_s
    else begin
      if lhs_width > 1
      then fprintf out_chan "assign %s[%d:%d] = %s;\n" lhs_net_name (hi-1) lo expr_s
      else fprintf out_chan "assign %s[%d] = %s;\n" lhs_net_name lo expr_s
    end;
    for i = 0 to lhs_width-1 do
      let driver_s = sprintf "%s[%d]" expr_s i in
      add_driver sym_info (lo+i) driver_s
    done
  end;
  if m.has_clock
  then begin
    fprintf out_chan "always @(posedge clk) begin\n";
    m.reg_alist
    |> List.iter ~f:
      begin fun (name, (_, next_expr, we_expr_opt)) ->
        let sym_info = find_net net_map name in
        if sym_info.kind <> Sym_reg
        then begin
          failwithf "LHS of '<=' is not reg: %s" name ()
        end;
        let next_expr, next_expr_width = check_expr net_map next_expr in
        check_width (Sym name, sym_info.width) (next_expr, next_expr_width);
        let next_expr_s = string_of_expr next_expr in
        match we_expr_opt with
        | Some we_expr ->
            let we_expr, _ = check_expr net_map we_expr in
            let we_expr_s = string_of_expr we_expr in
            fprintf out_chan "  if (%s) %s <= %s;\n" we_expr_s name next_expr_s
        | None ->
            fprintf out_chan "  %s <= %s;\n" name next_expr_s
      end;
    fprintf out_chan "end\n"
  end;
  List.iter m.instance_alist ~f:begin fun (inst_name, inst) ->
    let ports_s =
      inst.ports
      |> List.map ~f:begin fun (inst_name, _, expr) ->
        sprintf ".%s(%s)" inst_name (string_of_expr expr)
      end
      |> String.concat ~sep:", \n  "
    in
    if inst.parameters = []
    then fprintf out_chan "%s %s(\n  %s\n);\n" inst.class_name inst_name ports_s
    else begin
      let parameters_s =
        inst.parameters
        |> List.map ~f:(fun (a,b) -> sprintf ".%s(%d)" a b)
        |> String.concat ~sep:", "
      in
      fprintf out_chan "%s #(%s) %s(\n  %s\n);\n"
        inst.class_name parameters_s inst_name ports_s
    end;
    (* add driver to nets connected to output ports *)
    List.iter inst.ports ~f:begin fun (port_name, (port_width, direction), expr) ->
      match direction with
      | Input -> ()
      | Output ->
        begin match expr with
        | Sym net_name ->
            let sym_info = find_net net_map net_name in
            if port_width <> sym_info.width
            then begin
              failwithf "net width (%d) does not match port width (%d)"
                sym_info.width port_width ()
            end;
            for i = 0 to port_width-1 do
              let driver_s = sprintf "%s.%s[%d]" inst_name port_name i in
              add_driver sym_info i driver_s
            done
        | _ -> failwith "unimplemented"
        end
    end;
  end;
  fprintf out_chan "endmodule\n";
  Out_channel.close out_chan;
  (* raise if undriven [Sym_output] or [Sym_internal] nets exist *)
  Map.iteri net_map ~f:begin fun ~key ~data ->
    let net_name = key in
    let sym_info = data in
    match sym_info.kind with
    | Sym_input -> ()
    | Sym_output | Sym_internal ->
        for i = 0 to sym_info.width-1 do
          match sym_info.drivers.(i) with
          | [] -> failwithf "%s[%d] is not driven" net_name i ()
          | _ :: [] -> () (* exactly 1 driver, nothing to do *)
          | _ ->
              failwithf "%s[%d] has multiple drivers: %s" net_name i
                (String.concat sym_info.drivers.(i) ~sep:", ") ()
        done
    | Sym_reg -> ()
  end
